#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python

import numpy as np
import re
import os
import subprocess
import argparse
import sys
import math
class CCcore():
     
    def __init__(self,logFile):
        
        self.gaussianLog=logFile
        self.baseName='.'.join(self.gaussianLog.split('.')[:-1])
        print self.baseName
        self.shiftCoord=[]
        self.frequencies=[]
        self.originCoord=[]
        self.atomNum=[] 
        self.deltaCoord=[]
        self.cutPos=12
        self.deltaQ=0.1
        self.directories=[]
        self.level='b3lyp/3-21g'
        self.nproc='16'
        self.mem='8GB'
        self.spinCharge='0 1'
        self.orb1='h'
        self.orb2='h'
        self.transferInt=[]
        self.couplings=[]
        self.allOrbCoef1=[]
        self.allOrbCoef2=[]
        self.maxCoef1=[]
        self.maxCoef2=[]
        self.temp=298.15
    def readGaussianLog(self):
        logFile=open(self.gaussianLog,'r')
        switchSC=0
        switchO1=0
        switchO2=0
        for line in logFile:
            if 'Coord Atom Element:' in line:
                switchSC=switchSC+1
                i=0
            if switchSC==1 and re.search('^\s+\d+\s+\d+\s+\d+\s+-?\d+.\d+\s+',line):
                self.shiftCoord.append(map(float,line.split()))
            if switchSC>1 and re.search('^\s+\d+\s+\d+\s+\d+\s+-?\d+.\d+\s+',line):
                self.shiftCoord[i]=self.shiftCoord[i]+map(float,line.split()[3:])                   
                i=i+1
            if 'Frequencies ---' in line:
                self.frequencies=self.frequencies+line.split()[2:]
            if 'Standard orientation' in line:
                switchO1=1
                switchO2=0
                self.originCoord=[]
                self.atomNum=[]
            if '----------' in line and switchO1==1:
                switchO2=switchO2+1
            if switchO2==2 and len(line.split())==6:
                self.originCoord.append(float(line.split()[3]))
                self.originCoord.append(float(line.split()[4]))
                self.originCoord.append(float(line.split()[5])) 
                self.atomNum.append(line.split()[1])
        SC=np.array(self.shiftCoord)
        
        
        self.shiftCoord=SC.transpose()
        self.deltaCoord=[np.array([0.0]*len(self.originCoord))]
        #print self.deltaCoord
        #print len(self.shiftCoord)
        #print len(self.frequencies)
        #print self.originCoord
        #print self.atomNum
        for i in self.shiftCoord[3:]:
            self.deltaCoord.append(i)
        #print len(self.deltaCoord)
        for i in range(len(self.frequencies)+1):
            self.directories.append(self.baseName+'_'+str(int(self.deltaQ*1000)).zfill(4)+'_'+str(i).zfill(4)) 
        #print self.deltaCoord[0]
        
        
    def writeGaussianInput(self):
        for i in range(len(self.directories)):
            try: os.makedirs(self.directories[i]+'.cc')
            except OSError: pass
            os.chdir(self.directories[i]+'.cc')
            coords=''
            coordsFrg1=''
            coordsFrg2=''
            for atom in range(len(self.atomNum)):
                coords=coords + self.atomNum[atom]+'\t'+str(self.originCoord[atom*3+0]+self.deltaQ*self.deltaCoord[i][atom*3+0]) \
                                                  +'\t'+str(self.originCoord[atom*3+1]+self.deltaQ*self.deltaCoord[i][atom*3+1]) \
                                                  +'\t'+str(self.originCoord[atom*3+2]+self.deltaQ*self.deltaCoord[i][atom*3+2])+'\n'
                if atom < self.cutPos:
                    coordsFrg1=coords
                
                if atom >= self.cutPos:
                    coordsFrg2=coordsFrg2 + self.atomNum[atom]+'\t'+str(self.originCoord[atom*3+0]+self.deltaQ*self.deltaCoord[i][atom*3+0]) \
                                                  +'\t'+str(self.originCoord[atom*3+1]+self.deltaQ*self.deltaCoord[i][atom*3+1]) \
                                                  +'\t'+str(self.originCoord[atom*3+2]+self.deltaQ*self.deltaCoord[i][atom*3+2])+'\n'    

            #write dimer.gjf file
            dimer=open('dimer.gjf','w')
            dimer.write('%mem='+self.mem+'\n')
            dimer.write('%nprocshared='+self.nproc+'\n')
            dimer.write('# nosymm pop=nboread '+self.level+'\n \n')
            dimer.write(self.directories[i]+'\n \n')
            dimer.write(self.spinCharge+'\n')
            dimer.write(coords+' \n')
            dimer.write('$NBO SAO=w13 FAO=w14 $END\n \n \n')
            dimer.close()
            monomer=open('frag1.gjf','w')
            monomer.write('%mem='+self.mem+'\n')
            monomer.write('%nprocshared='+self.nproc+'\n')
            monomer.write('# nosymm punch(MO) '+self.level+'\n \n')
            monomer.write(self.directories[i]+'\n \n')
            monomer.write(self.spinCharge+'\n')
            monomer.write(coordsFrg1+' \n')
            monomer.write('\n \n')
            monomer.close()
            monomer=open('frag2.gjf','w')
            monomer.write('%mem='+self.mem+'\n')
            monomer.write('%nprocshared='+self.nproc+'\n')
            monomer.write('# nosymm punch(MO)'+self.level+'\n \n')
            monomer.write(self.directories[i]+'\n \n')
            monomer.write(self.spinCharge+'\n')
            monomer.write(coordsFrg2+' \n')
            monomer.write('\n \n')
            monomer.close()
            data=open('data.in','w')
            data.write('FILE.13\n')
            data.write('FILE.14\n')
            data.write('frag1.mo\n')
            data.write('frag2.mo\n')
            data.close()
            script1=open('runG09.sh','w')
            script1.write('g09 dimer.gjf\n')
            script1.write('g09 frag1.gjf\n')
            script1.write('mv fort.7 frag1.mo\n')
            script1.write('g09 frag2.gjf\n')
            script1.write('mv fort.7 frag2.mo\n')
            script1.close()
            script2=open('runCouple.sh','w')
            script2.write('echo BasisFunctions1= `grep "basis functions" frag1.log | awk \'{print $1}\'` > param.in\n')
            script2.write('echo BasisFunctions1= `grep "basis functions" frag2.log | awk \'{print $1}\'` >> param.in\n')
            orb1='0'
            orb2='0'
            if self.orb1.lower()[0]=='h':
                if self.orb1.lower().strip('h'):
                    orb1=self.orb1.lower().strip('h')
                script2.write('mo=`grep "alpha electrons" frag1.log | awk \'{print $1-'+orb1+'}\'`\n')
                script2.write('echo NumeroOrbital1= $mo,$mo >> param.in\n')
            elif self.orb1.lower()[0]=='l':
                if self.orb1.lower().strip('l'):
                    orb1=self.orb1.lower().strip('l')
                script2.write('mo=`grep "alpha electrons" frag1.log | awk \'{print $1+1+'+orb1+'}\'`\n')
                script2.write('echo NumeroOrbital1= $mo,$mo >> param.in\n')
            else:
                script2.write('echo NumeroOrbital1= '+self.orb1+','+self.orb1+' >> param.in\n')
                
            if self.orb2.lower()[0]=='h':
                if self.orb2.lower().strip('h'):
                    orb2=self.orb2.lower().strip('h')
                script2.write('mo=`grep "alpha electrons" frag2.log | awk \'{print $1-'+orb2+'}\'`\n')
                script2.write('echo NumeroOrbital2= $mo,$mo >> param.in\n')
            elif self.orb2.lower()[0]=='l':
                if self.orb2.lower().strip('l'):
                    orb2=self.orb2.lower().strip('l')
                script2.write('mo=`grep "alpha electrons" frag2.log | awk \'{print $1+1+'+orb2+'}\'`\n')
                script2.write('echo NumeroOrbital2= $mo,$mo >> param.in\n')
            else:
                script2.write('echo NumeroOrbital2= '+self.orb2+','+self.orb2+' >> param.in\n')
            script2.write('ccc')
            os.chdir("..")
            
    def calculateCC(self):
         print 'deltaQ is set to '+str(self.deltaQ)
         print "{:<6s} {:<10s} {:<12s} {:<12s} {:<12s} {:<12s} {:<12s} {:<6s} {:<12s} {:<12s}".format('modes','frequency','t (eV)','delta t (eV)','coupling','max_coef1','max_coef2','corr','dt correct' ,'coupl. correct')
         maxIndex1=-1
         maxIndex2=-1
         corrections=[]
         for modes in range(len(self.directories)):
            os.chdir(self.directories[modes]+'.cc') 
            #print self.directories[i]
            try:
                os.remove('t-values.out')
            except OSError: pass
            try:
                devnull = open(os.devnull, 'wb')
                runCCC=subprocess.Popen(['sh','runCouple.sh'], stderr=devnull)
                runCCC.wait()  
                result=open('t-values.out','r')
            except IOError:
                runG09=subprocess.Popen(['sh','runG09.sh'])
                runG09.wait()   
                runCCC=subprocess.Popen(['sh','runCouple.sh'])
                runCCC.wait()   
                result=open('t-values.out','r')
            couplings=[]
            transferInt=[]
            for line in result:
                transferInt.append(float(line.split('=')[1].split()[0]))
            self.transferInt.append(transferInt)
            #check orbital phase
            #read the orbital
            param=open('param.in','r')
            lines=param.readlines()
            orb1=lines[2].split(',')[1].strip()
            orb2=lines[3].split(',')[1].strip()
            orbFile1=open('frag1.mo','r')
            switch=0
            orbCoef1=[]
            for line in orbFile1:
                if orb1+' Alpha MO' in line:
                    switch=1
                    continue
                elif 'Alpha MO' in line:
                    switch=0
                if switch==1:
                    orbCoef1=orbCoef1+map(float,filter(None,[line.strip('\n').replace('D','E')[i:i+15] for i in range(0,60,15)]))
            self.allOrbCoef1.append(orbCoef1)
            if maxIndex1==-1:
                maxIndex1=self.allOrbCoef1[0].index(max(self.allOrbCoef1[0])) 
            orbFile2=open('frag2.mo','r')
            switch=0
            orbCoef2=[]
            for line in orbFile2:
                if orb2+' Alpha MO' in line:
                    switch=1
                    continue
                elif 'Alpha MO' in line:
                    switch=0
                if switch==1:
                    orbCoef2=orbCoef2+map(float,filter(None,[line.strip('\n').replace('D','E')[i:i+15] for i in range(0,60,15)]))
            self.allOrbCoef2.append(orbCoef2)
            if maxIndex2==-1:
                maxIndex2=self.allOrbCoef2[0].index(max(self.allOrbCoef2[0]))   
            if orbCoef2[maxIndex2]*orbCoef1[maxIndex1] > 0:
                correction=1
            else:
                correction=-1         
            corrections.append(correction)
            self.couplings.append((transferInt[0] -self.transferInt[0][0])/self.deltaQ)
            print "{:<6d} {:<10s} {:< 12.7f} {:< 12.7f} {:< 12.7f} {:< 12.7f} {:< 12.7f} {:< 6d} {:< 12.7f} {:< 12.7f}".format(modes,(['0']+self.frequencies)[modes],transferInt[0],transferInt[0]-self.transferInt[0][0],(transferInt[0] -self.transferInt[0][0])/self.deltaQ,orbCoef1[maxIndex1],orbCoef2[maxIndex2],correction,transferInt[0]*correction-self.transferInt[0][0],(transferInt[0]*correction -self.transferInt[0][0])/self.deltaQ)
            os.chdir("..")
         maxIndex1=self.allOrbCoef1[0].index(max(self.allOrbCoef1[0]))
         for i in self.allOrbCoef1:
            self.maxCoef1.append(i[maxIndex1])
         maxIndex2=self.allOrbCoef2[0].index(max(self.allOrbCoef2[0]))
         for i in self.allOrbCoef2:
            self.maxCoef2.append(i[maxIndex2])
         #print "{:<6s} {:<10s} {:<12s} {:<12s} {:<12s}".format('modes','frequency','max_coef1','max_coef2', 'corrections')   
         #for i in range(len(self.frequencies)+1):
         #    print "{:<6d} {:<10s} {:< 12.7f} {:< 12.7f} {:<12d}".format(i,(['0']+self.frequencies)[i],self.maxCoef1[i],self.maxCoef2[i], corrections[i])   
        
         #G2=0
         #L=0
         #sigma=0
         #for i in range(1,len(self.couplings)):
         #    G2=G2+self.couplings[i]**2/2
         #    L=L+self.couplings[i]**2/(2*4.135667516E-15*29979245800*float(self.frequencies[i-1]))
         #    sigma=sigma+self.couplings[i]**2/2*math.cosh((4.135667516E-15*29979245800*float(self.frequencies[i-1]))/(2*8.6173324E-5*self.temp))
         #G=G2**0.5
         #print "G:  "+str(G)+" ;  L:  "+str(L)+" ;  Sigma2:  "+str(sigma)
parser=argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,description="A script used with Gaussian09 to auto calculate coupling constant.")
parser.add_argument("-l","--level",help="Set the theoretical level. Default: read from gjf file. If not found, use b3lyp/6-31g(d)",  type=str,default='b3lyp/6-31g(d)')
parser.add_argument("-m","--memory",help="Set the memory used by Gaussian (with unit).  Default: 500MB",  type=str,default='500MB')  
parser.add_argument("-n","--nproc",help="Set the number of cores used by Gaussian. Default: 8",  type=str,default='8') 
parser.add_argument("-d","--delta",help="Set the delta of Q. Default: 0.1", default=0.1,type=float)
parser.add_argument("-f","--fragment",help="Set the number of atoms of fragment 1. ", default=None, type=int)
parser.add_argument("-t","--temperature",help="Set the temerature to calculate sigma. ", default=298.15, type=float)
parser.add_argument("-o","--orbital",help="Set the orbital for fragment 1 and fragment 2, seperate by comma. Can use relative or absolute value. e.g. h1 stands for HOMO-1, l2 stands for LUMO+2, 134 stands for the 134th orbital. Default: H,H", default="h,h", type=str)
parser.add_argument("-c","--compute",help="set this to 0 for generate input file and directory only", type=str,default='1') 
parser.add_argument("gaussianOutput",nargs=1)
args=parser.parse_args()
if not args.fragment:
    print 'Must input number of atoms of fragment 1 (use -f option)'
    sys.exit()
ccc=CCcore(args.gaussianOutput[0])
ccc.temp=args.temperature
ccc.cutPos=args.fragment
ccc.deltaQ=args.delta
ccc.level=args.level
ccc.mem=args.memory
ccc.orb1=args.orbital.split(',')[0]
ccc.orb2=args.orbital.split(',')[1]
ccc.readGaussianLog()
ccc.writeGaussianInput()
if args.compute!='0':
    ccc.calculateCC()
    